package org.lionsoul.websnail.downloader.webdriver;

import org.apache.logging.log4j.LogManager;
import org.apache.logging.log4j.Logger;
import org.lionsoul.websnail.Spider;
import org.lionsoul.websnail.downloader.Downloader;
import org.lionsoul.websnail.downloader.WebDriverProcessor;
import org.openqa.selenium.By;
import org.openqa.selenium.TimeoutException;
import org.openqa.selenium.WebElement;
import org.openqa.selenium.support.ui.WebDriverWait;

import java.time.Duration;

public class WebDriverDownloader implements Downloader
{
    private static final Logger logger = LogManager.getLogger(WebDriverDownloader.class);
    private final WebDriverBuilder webDriverBuilder;
    private WebDriverProcessor driverProcessor;

    public WebDriverDownloader(WebDriverBuilder builder, WebDriverProcessor driverProcessor)
    {
        this.webDriverBuilder = builder;
        this.driverProcessor = driverProcessor;
        // @BugFix: 设置 chrome driver 路径，兼容某些系统获取不到 chrome driver 的情况
        String chromeDriver = System.getenv("CHROME_DRIVER");
        if (null != chromeDriver) {
            System.setProperty("webdriver.chrome.driver", chromeDriver);
        }
    }

    public WebDriverDownloader(WebDriverBuilder builder)
    {
        this(builder, null);
    }

    public WebDriverProcessor getDriverProcessor()
    {
        return driverProcessor;
    }

    public void setDriverProcessor(WebDriverProcessor driverProcessor)
    {
        this.driverProcessor = driverProcessor;
    }

    @Override
    public String process(Spider spider, String url)
    {
        SnailWebDriver driver = webDriverBuilder.build(driverProcessor);
        // if the WebDriver pool is full, and no available WebDriver
        // then we should rebuild later
        while (null == driver) {
            try {
                Thread.sleep(spider.getOptions().getWorkerBusySleepTime());
            } catch (InterruptedException e) {
                logger.info("The WebDriver pool is full and all WebDriver is busing, rebuild later...");
            }
            driver = webDriverBuilder.build(driverProcessor);
        }
        driver.getWebDriver().get(url);

        // 设置页面加载完成条件
        if (null != driver.getCondition()) {
            try {
                WebDriverWait wait = new WebDriverWait(driver.getWebDriver(), Duration.ofSeconds(5), Duration.ofSeconds(1));
                wait.until(driver.getCondition());
            } catch (TimeoutException e) {
                logger.warn(e.getMessage());
            }
        }

        try {
            WebElement webElement = driver.getWebDriver().findElement(By.tagName("html"));
            return webElement.getAttribute("outerHTML");
        } catch (Exception e) {
            return null;
        } finally {
            driver.setState(SnailWebDriver.IDLE);
        }
    }

    public WebDriverBuilder getWebDriverBuilder()
    {
        return webDriverBuilder;
    }

    @Override
    public void shutdown()
    {
        webDriverBuilder.close();
    }
}
